# -*- coding: utf-8 -*-
import time
import os
import sys
from os.path import dirname

# -----------------------------
# 🔧 参数定义（在这里直接改即可）
# -----------------------------
BASE_DIR = dirname(os.path.abspath(__file__))       # 项目根目录

# 测试图片路径
TEST_DATA = os.path.join(BASE_DIR, 'my_dataset', 'orl_faces_processed','test', 's1_1.jpg')

# 特征索引文件路径
INDEX_FILE = os.path.join(BASE_DIR, 'index', 'my_train.h5')

# 数据库名称（Milvus）
DB_NAME = 'image_retrieval_new'

# 检索引擎，可选："numpy" 或 "milvus"
ENGINE = 'milvus'
# -----------------------------

# 添加模块路径
sys.path.append(BASE_DIR)

from service.vggnet import VGGNet
from service.numpy_retrieval import NumpyRetrieval
from service.milvus_retrieval_new import MilvusRetrieval


class RetrievalEngine(object):
    """检索引擎封装类"""

    def __init__(self, index_file, db_name):
        self.index_file = index_file
        self.db_name = db_name
        self.numpy_r = self.milvus_r = None

    def get_method(self, m_name):
        """根据引擎名获取对应处理方法"""
        m_name = f"{m_name}_handler"
        return getattr(self, m_name, self.default_handler)

    def numpy_handler(self, query_vector, req_id=None):
        """使用 Numpy 检索"""
        if self.numpy_r is None:
            self.numpy_r = NumpyRetrieval(self.index_file)
        return self.numpy_r.retrieve(query_vector)

    def milvus_handler(self, query_vector, req_id=None):
        """使用 Milvus 检索"""
        if self.milvus_r is None:
            self.milvus_r = MilvusRetrieval(self.db_name, self.index_file)
        return self.milvus_r.retrieve(query_vector)

    def default_handler(self, query_vector, req_id=None):
        """默认处理函数（无结果）"""
        return []


if __name__ == '__main__':
    start_time = time.time()

    # -----------------------------
    # 🧠 步骤 1：提取图片特征
    # -----------------------------
    print(f"正在提取特征: {TEST_DATA}")
    model = VGGNet()
    query_vector = model.vgg_extract_feat(TEST_DATA)

    # -----------------------------
    # 🔍 步骤 2：执行检索
    # -----------------------------
    print(f"使用引擎: {ENGINE}")
    re_engine = RetrievalEngine(INDEX_FILE, DB_NAME)
    result = re_engine.get_method(ENGINE)(query_vector, None)

    # -----------------------------
    # ⏱ 输出结果与耗时
    # -----------------------------
    end_time = time.time()
    print(f"\n🔹 检索结果: {result}")
    print(f"⏱ 总耗时: {end_time - start_time:.2f} 秒")
